#include "rsa.h"
#include "string.h"
#include "winbase.h"

class prime_factory
{
  unsigned np;
  unsigned *pl;
  public:
  prime_factory();
  ~prime_factory();
  vlong find_prime( vlong & start );
};

// prime factory implementation

static int is_probable_prime( const vlong &p )
{
  // Test based on Fermats theorem a**(p-1) = 1 mod p for prime p
  // For 1000 bit numbers this can take quite a while
  const int rep = 4;
  const unsigned any[rep] = { 2,3,5,7 };
  for ( unsigned i=0; i<rep; i+=1 )
    if ( modexp( any[i], p-1, p ) != 1 )
      return 0;
  return 1;
}

prime_factory::prime_factory()
{
  np = 0;
  unsigned NP = 200;
  pl = new unsigned[NP];

  // Initialise pl
  unsigned SS = 8*NP; // Rough estimate to fill pl
  char * b = new char[SS+1]; // one extra to stop search
  for (unsigned i=0;i<=SS;i+=1) b[i] = 1;
  unsigned p = 2;
  while (1)
  {
    // skip composites
    while ( b[p] == 0 ) p += 1;
    if ( p == SS ) break;
    pl[np] = p;
    np += 1;
    if ( np == NP ) break;
    // cross off multiples
    unsigned c = p*2;
    while ( c < SS )
    {
      b[c] = 0;
      c += p;
    }
    p += 1;
  }
  delete [] b;
}

prime_factory::~prime_factory()
{
  delete [] pl;
}

vlong prime_factory::find_prime( vlong & start )
{
  unsigned SS = 1000; // should be enough unless we are unlucky
  char * b = new char[SS]; // bitset of candidate primes
  unsigned tested = 0;
  while (1)
  {
    unsigned i;
    for (i=0;i<SS;i+=1)
      b[i] = 1;
    for (i=0;i<np;i+=1)
    {
      unsigned p = pl[i];
      unsigned r = start % (vlong)p; // not as fast as it should be - could do with special routine
      if (r) r = p - r;
      // cross off multiples of p
      while ( r < SS )
      {
        b[r] = 0;
        r += p;
      }
    }
    // now test candidates
    for (i=0;i<SS;i+=1)
    {
      if ( b[i] )
      {
        tested += 1;
        if ( is_probable_prime(start) )
          return start;
      }
      start += 1;
    }
  }
  delete [] b;
}

void private_key::create( const char * r1, const char * r2 )
{
  // Choose primes
  {
    prime_factory pf;
    p = pf.find_prime( from_str(r1) );
    q = pf.find_prime( from_str(r2) );
    if ( p > q ) { vlong tmp = p; p = q; q = tmp; }
  }
  // Calculate public key
  {
    m = p*q;
    e = 50001; // must be odd since p-1 and q-1 are even
    while ( gcd(p-1,e) != 1 || gcd(q-1,e) != 1 ) e += 2;
  }
}

vlong public_key::encrypt( const vlong& plain )
{
  return modexp( plain, e, m );
}

vlong private_key::decrypt( const vlong& cipher )
{
  // Calculate values for performing decryption
  // These could be cached, but the calculation is quite fast
  vlong d = modinv( e, (p-1)*(q-1) );
  vlong u = modinv( p, q );
  vlong dp = d % (p-1);
  vlong dq = d % (q-1);

  // Apply chinese remainder theorem
  vlong a = modexp( cipher % p, dp, p );
  vlong b = modexp( cipher % q, dq, q );
  if ( b < a ) b += q;
  return a + p * ( ((b-a)*u) % q );
}

vlong public_key::encrypt( const char* plain, const char * exponant, const char * modulus, int keylength )
{

	int slen = 0;
	char * sp = (char *)plain;
	while(plain[slen] != 0) slen++;
	int npad = keylength - 3 - slen;
	char * s = new char[keylength+1];
	char * p = s;

	*p++ = 0;
	*p++ = 2;
	for(int i = 2; i < npad + 2; i++)
	{
		*p++ = rand() % 255 + 1;
	}
	*p++ = 0;
	while(*sp != 0)	*p++ = *sp++;
	*p = 0;

	e = from_hexstr(exponant);
	m = from_hexstr(modulus);
	vlong res;
	res = from_str(s,keylength);
	delete(s);
	res = modexp( res, e, m );
	return res;
}

//#include "CStr.h"

vlong public_key::pkcs1_encrypt( AStr message, const char * publickey, const char * modulus, int KeyLenth, bool isPublickey )
{
	//Trace(L"pkcs1 In\r\n");
	int nb_pad;
	int ilen=message.length();
	char * Output = (char *)malloc(KeyLenth+ilen+4);
	if ( Output == NULL ) return 0;
	unsigned char * p = (unsigned char *)Output;
	vlong i,o,e,n;
	
	nb_pad = KeyLenth - 3 - ilen;
	*p++ = 0;
	if (isPublickey)
	{
		*p++ = 0x02;
		
		srand(GetTickCount());
		while(nb_pad-- > 0)
			*p++ =  rand() % 254  + 1;
	}
	else
	{
		*p++ = 0x01;
		while(nb_pad-- > 0)
			*p++ = 0xFF;
	}

	*p++ = 0;
	memcpy(p,(char*)message,ilen);
	p+=ilen;
	*p = 0;
	i = from_str(Output,KeyLenth);
	e = from_hexstr(publickey);
	n = from_hexstr(modulus);

	return modexp(i,e,n);
}


/*
#include "stdio.h"

int main()
{
	vlong x,y,z;
	//int h;
	//B2AEA60BFE3127A3B7DF8C1DB454867E25B63F1D0C101C93337CC3EAE6FAA1E27CAB5BC5E1A3CC0EBA78038002EB08DEEF7C288EDCE14B8D3731D49A52BBD7758BC8F8A856B2126588093D8F0ADC141151D72B930C8D8486432D7DBF0BE298B75AF9BD24EB4352A8426AEBA5BA317945C8DEC790CFBDECB20DDF484A49B42165010001
	x=from_hexstr("B2AEA60BFE3127A3B7DF8C1DB454867E25B63F1D0C101C93337CC3EAE6FAA1E27CAB5BC5E1A3CC0EBA78038002EB08DEEF7C288EDCE14B8D3731D49A52BBD7758BC8F8A856B2126588093D8F0ADC141151D72B930C8D8486432D7DBF0BE298B75AF9BD24EB4352A8426AEBA5BA317945C8DEC790CFBDECB20DDF484A49B42165");
	y=from_hexstr("010001");
	z=from_str("fetion.com.cn:j11w00d");
	private_key r;
	r.e = y;
	r.m = x;
	x = r.encrypt(z);
	char * s = to_hexstr(x);
	printf("%s\n",s);
	delete(s);
	return 0;
}
*/